
# 战斗、采集和保存数据
# 战斗暂时用人手操，后面换成神经网络预测+键盘模拟
# 数据包括：战斗画面(压缩)、双方血量及对应奖励信号；保存方式？

import cv2
import mss
import time
import threading
import numpy as np
import pydirectinput
from pynput import keyboard
from Activate_Window import activate_window
from config import (FPS, DURATION, KNIGHT_HP_THRESHOLD,
                    CAPTURE_WINDOW, KNIGHT_HP_WINDOW, BOSS_HP_WINDOW,
                    ALLOWED_KEYS)

current_held_keys = set()          # 当前被按下的合法键
new_keys_this_frame = []           # 本帧新按下的合法键
keys_lock = threading.Lock()       # 保护共享变量
def on_press(key):
    global current_held_keys, new_keys_this_frame
    with keys_lock:
        try:
            key_name = key.char.upper()
        except AttributeError:
            if key == keyboard.Key.space:
                key_name = 'SPACE'
            elif key == keyboard.Key.esc:
                key_name = 'ESC'  # 保留 ESC 用于退出
            else:
                return  # 忽略其他特殊键

        # 只处理我们关心的键
        if key_name not in ALLOWED_KEYS:
            return

        # 如果是新按下的，记录为“本帧新动作”
        if key_name not in current_held_keys:
            current_held_keys.add(key_name)
            new_keys_this_frame.append(key_name)  # 记录新动作
def on_release(key):
    global current_held_keys
    with keys_lock:
        try:
            key_name = key.char.upper()
        except AttributeError:
            if key == keyboard.Key.space:
                key_name = 'SPACE'
            elif key == keyboard.Key.esc:
                key_name = 'ESC'
            else:
                return  # 忽略其他特殊键

        # 只处理我们关心的键
        if key_name in ALLOWED_KEYS and key_name in current_held_keys:
            current_held_keys.remove(key_name)

# 原始画面截图
def capture_frame(sct):
    img = np.array(sct.grab(CAPTURE_WINDOW))    # BGRA (H, W, 4)
    return img[:, :, :3]                        # BGR  (H, W, 3)

# 从原图提取小骑士、BOSS血量图
def extract_hp_bar(frame,window):
    y = window["top"]
    x = window["left"]
    dy = window["height"]
    dx = window["width"]
    return frame[y:y+dy, x:x+dx]

# 读一次小骑士血量参照图
KNIGHT_REF_IMG = cv2.imread("../knight_ref_hp_laptop.png")

# 识别小骑士血量图为整数数值信息
def detect_knight_hp_change(cmp_img,threshold=KNIGHT_HP_THRESHOLD):
    match = (cmp_img == KNIGHT_REF_IMG).all(axis=2)[0]  # shape: (570,)
    diff = np.diff(match.astype(np.int8), prepend=0, append=0)
    starts = np.where(diff == 1)[0]  # True 段开始
    ends = np.where(diff == -1)[0]  # True 段结束
    change = np.sum(ends - starts >= threshold)
    return change

# 识别boss血量图为数值信息
def detect_boss_hp_percentage(hp_bar_img):
    # [5, 0, 96]可用get_hp_BGR获得
    if not np.array_equal(hp_bar_img[0, 0], [5, 0, 96]):
        return 0.0
    w = hp_bar_img.shape[1]
    diff = np.diff(hp_bar_img[0, :, :].astype(np.int16), axis=0)
    candidates = np.where(np.abs(diff).max(axis=1) > 1)[0]
    if len(candidates) > 0:
        return candidates[0] / w
    else:
        return 1.0

def main():
    activate_window("Hollow Knight")

    pydirectinput.press("w")
    time.sleep(2)
    pydirectinput.press("w")
    time.sleep(1)
    pydirectinput.press("space")
    time.sleep(3)

    with mss.mss() as sct:
        start_time = time.time()

        count = 0                   # 轮次唯一标识
        prev_knight = 10         # 小骑士上轮生命值
        prev_knight_loss = 0     # 小骑士上轮受伤值

        prev_boss = 0

        while True:
            next_frame_time = start_time + (count + 1) / FPS
            current_time = time.time()
            if current_time < next_frame_time:
                time.sleep(next_frame_time - current_time)

            frame = capture_frame(sct)
            count += 1

            # 🔑 关键：先获取数据，再清空 new_keys，避免竞争
            with keys_lock:
                # 1. 获取本帧新按下的键
                newly_pressed = new_keys_this_frame.copy()

                # 2. 获取当前所有被按下的合法键（包括新按 + 长按）
                currently_held = current_held_keys.copy()

                # 3. 清空 new_keys，为下一帧准备
                new_keys_this_frame.clear()

            # 4. 计算“持续按住”的键：当前按着的 - 本帧新按下的
            held_without_new = currently_held - set(newly_pressed)

            # 🖨️ 输出
            if newly_pressed:
                new_str = "、".join(newly_pressed)
            else:
                new_str = "无"

            if held_without_new:
                held_str = "、".join(held_without_new)
            else:
                held_str = "无"

            print(f"第{count}步：新按下 {new_str}；持续按住 {held_str}")

            # 分割frame为三个部分
            # ①boss血量
            boss_hp_bar = extract_hp_bar(frame,BOSS_HP_WINDOW)
            current_boss = detect_boss_hp_percentage(boss_hp_bar)
            # ②小骑士血量
            knight_hp_bar = extract_hp_bar(frame,KNIGHT_HP_WINDOW)
            current_knight_loss = detect_knight_hp_change(knight_hp_bar,KNIGHT_HP_THRESHOLD)
            current_knight = prev_knight - (current_knight_loss if current_knight_loss != prev_knight_loss else 0)
            if current_knight_loss != prev_knight_loss or prev_knight_loss == 0:
                prev_knight_loss = current_knight_loss
            # ③战斗画面


            reward = 0

            is_knight_dead = False
            is_boss_dead = False
            is_error_happen = False

            # 优先分析小骑士，因为其血量提供的信息更加真实；boss血量在最后不确定
            if current_knight == prev_knight: # 无伤
                reward += 1 # 微微鼓励无伤
            elif current_knight == 0:  # 角色死亡，采样结束，先不break等待boss血量结算
                reward -= 10
                # print(f"第{count}步 : 受伤{prev_knight-current_knight}血！角色死亡！reward:{reward}")
                is_knight_dead = True
            elif current_knight < prev_knight:
                reward = -5 # 受伤惩罚比造成伤害更大，先学躲技能再学砍人
                # print(f"第{count}步 : 受伤{prev_knight-current_knight}血！reward:{reward}")
            # 这里没有考虑小骑士加血的情况，即在当前假设下血量是单调递减的
            else:
                reward -= 999
                is_error_happen = True
                # print("小骑士血量出错！")

            if current_boss == prev_boss: # 开头还没攻击到、中间没攻击到
                pass
            # 单boss单阶段，这就是击败boss的瞬间
            elif current_boss == 0 and prev_boss != 0:
                reward += 10
                # print(f"第{count}步 : 造成伤害{100*(prev_boss):.1f} 击杀BOSS！ reward:{reward}")
                is_boss_dead = True
            elif current_boss < prev_boss:
                reward += 1
                # print(f"第{count}步 : 造成伤害{100*(prev_boss-current_boss):.1f} reward:{reward}")
            elif prev_boss == 0 and current_boss != 0:
                reward += 5
                # print(f"第{count}步 : 造成伤害{100*(1-current_boss):.1f} 初次攻击！reward:{reward}")
            # 这里后面可以尝试处理的逻辑是：boss血量恢复了，部分回复或者回复到1
            # 现在就只打单boss单阶段是没问题的
            else:
                reward -= 999
                is_error_happen = True
                # print("BOSS血量出错！")

            if is_knight_dead or is_boss_dead or is_error_happen:
                break
            if (time.time() - start_time) >= DURATION:
                # print(f"{time.time() - start_time:.1f}s : ⏹️ 采样结束(未分出胜负，需要加长时间)")
                break

            prev_knight = current_knight
            prev_boss = current_boss

        total_time = time.time() - start_time
        actual_fps = count / total_time  # 帧数 / 总时间

        print(f"\n📊 采样统计:")
        print(f"   目标时长: {DURATION:.1f}s")
        print(f"   实际耗时: {total_time:.2f}s")
        print(f"   总帧数: {count}")
        print(f"   实际平均采样率: {actual_fps:.2f} FPS")


if __name__ == '__main__':
    listener = keyboard.Listener(on_press=on_press, on_release=on_release)
    listener.start()
    main()
    listener.stop()  # 确保监听器停止
    listener.join()  # 等待线程结束


